import torch
from diffusers import HunyuanVideoTransformer3DModel

@torch.no_grad()
def main():
    device = "cuda"
    dtype = torch.bfloat16
    batch_size = 1
    num_channels = 4
    num_frames = 2
    height = 4
    width = 4
    sequence_length = 8
    
    transformer = HunyuanVideoTransformer3DModel(
        in_channels=4,
        out_channels=4,
        num_attention_heads=2,
        attention_head_dim=10,
        num_layers=2,
        num_single_layers=2,
        num_refiner_layers=1,
        patch_size=1,
        patch_size_t=1,
        guidance_embeds=True,
        text_embed_dim=16,
        pooled_projection_dim=8,
        rope_axes_dim=(2, 4, 4),
    ).to(device=device, dtype=dtype)

    hidden_states = torch.randn(batch_size, num_channels, num_frames, height, width, dtype=dtype, device=device)
    timestep = torch.randint(0, 1000, (batch_size,), dtype=torch.long, device=device)
    encoder_hidden_states = torch.randn(batch_size, sequence_length, 16, dtype=dtype, device=device)
    encoder_attention_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
    encoder_attention_mask[:, :sequence_length - 2] = True
    pooled_projections = torch.randn(batch_size, 8, dtype=dtype, device=device)
    guidance = torch.randint(0, 1000, (batch_size,), dtype=torch.long, device=device)

    transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
    output = transformer(
        hidden_states=hidden_states,
        timestep=timestep,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_attention_mask,
        pooled_projections=pooled_projections,
        guidance=guidance,
        return_dict=False,
    )[0]
    print(output.shape)


if __name__ == "__main__":
    main()